#include "include/LPAS.h"

void LPAS::assign_flow(const int id, const int priority, bool is_scheduled) {
	flows[id].priority     = priority;
    flows[id].is_scheduled = is_scheduled;
    // if (move_from_unassign) {
		erase_unassign(id);
    // }
}

void LPAS::run() {
    int priority = flow_num;
    clear_assign();
    clear_ep();
    clear_unassign();
    // 按照关键性升序，相对截止时间降序，周期降序排序
    // 排前面的优先级低（分配到的优先级数字大）
    sort(flows + 1, flows + flow_num + 1, [&](flow& a, flow& b) {
        if (a.critical == b.critical) {
            if (a.relative_dl == b.relative_dl) {
                return a.period > b.period;
            }
            return a.relative_dl > b.relative_dl;
        }
        return a.critical < b.critical;
    });
    for (int i = 1; i <= flow_num; i++) {
        add_unassign(i);
        flows[i].is_scheduled = false;
    }
    while (!assign_all()) {
        bool is_schedule = false;
        int  schedule_id = -1;
        for (auto cur_id : all_unassign_flow_set) {
			add_ep(cur_id);
            if (WCRT(flows[cur_id])) {
                assign_flow(cur_id, priority);
                is_schedule = true;
                schedule_id = cur_id;
                break;
            }
            erase_ep(cur_id);
        }
		// printf("prio:%d, schedule:%d, id:%d\n",priority,is_schedule,flows[schedule_id].id);
        if (!is_schedule) {
            auto now = all_unassign_flow_set.begin();
            add_ep(*now);
            erase_unassign(*now);
        } else {
            for (auto cur_id : all_ep_flow_set) {
                assign_flow(cur_id, priority, WCRT(flows[cur_id]));
            }
            for (auto cur_id : all_ep_flow_set)
                add_assign(cur_id);
            add_assign(schedule_id);
            clear_ep();
			priority--;
        }
    }

	for (auto cur_id : all_ep_flow_set) {
		// printf("last prio:%d, schedule:%d, id:%d\n",priority,is_schedule,flows[cur_id].id);
		assign_flow(cur_id, priority, WCRT(flows[cur_id]));
	}
    for (int i = 1;i <= flow_num;i++)
        flows[i].old_priority = flows[i].priority;
    // LPR();
}

void LPAS::separate_ep_flows() {
	sort(flows + 1, flows + flow_num + 1, [&](flow& a, flow& b) {
        if (a.priority == b.priority) {
            return a.critical < b.critical;
        }
        return a.priority > b.priority;
    });
	clear_assign();
    clear_ep();
    clear_unassign();
	for (int i = 1; i <= flow_num; i++)
        add_unassign(i);
    for (int i = 1; i <= flow_num; i++) {
    	int j = i;
		while (j + 1 <= flow_num && flows[j + 1].priority == flows[i].priority) {
			j++;
		}
		if (j == i) {
			flows[i].new_priority = flows[i].priority;
			add_assign(i);
			erase_unassign(i);
			continue;
		}
		int len = j - i + 1;
        if (len > 25) {
            for (int k = i;k <= flow_num;k++)
                flows[k].new_priority = flows[k].priority;
            break;
        }
        if (len > 10) {
            for (int k = i;k <= j;k++) {
                flows[k].new_priority = flows[k].priority;
                add_assign(k);
                erase_unassign(k);
            }
            i = j;
            continue;
        }
		u64 situation = (1L << (len - 1));
		int hi_scheduled = 0, lo_scheduled = 0, best_situation = 0;
        cp_set(assign_flow_set, temp_assign_set);
        cp_set(unassign_flow_set, temp_unassign_set);
		for (u64 k = 0; k < situation; k++) {
			int new_hi_scheduled = 0, new_lo_scheduled = 0;
			for (int sep = 0;sep <= len - 1;sep++) {
                clear_ep();
				int left = i + sep, right = left;
				while (right <= j) {
                    add_ep(right);
					erase_unassign(right);
					// printf("k:%d, l:%d, r:%d, sep:%d, res:%d\n", k, left, right, sep, (k & (1L << sep)));
					if (sep == len - 1 || (k & (1L << sep))) break;
					sep ++;
					right ++;
				}
				for (int l = left;l <= right;l++) {
					if (WCRT(flows[l])) {
						flows[l].critical == 1 ? new_hi_scheduled ++ : new_lo_scheduled ++;
					}
				}
				for (int l = left;l <= right;l++)
					add_assign(l);
			}
			if (new_hi_scheduled > hi_scheduled || (new_hi_scheduled == hi_scheduled && new_lo_scheduled > lo_scheduled)) {
				// printf("separate:%d prio:%d, newhi:%d, newlo:%d, oldhi:%d, oldlo:%d\n", k, flows[i].priority, new_hi_scheduled, new_lo_scheduled, hi_scheduled, lo_scheduled);
                hi_scheduled = new_hi_scheduled;
				lo_scheduled = new_lo_scheduled;
				best_situation = k;
			}
			cp_set(temp_assign_set, assign_flow_set);
            cp_set(temp_unassign_set, unassign_flow_set);
		}
		// 执行最优分割
		int prio = flows[i].priority;
		for (int sep = 0;sep <= len - 1;sep++) {
			clear_ep();
			int left = i + sep, right = left;
			while (right <= j) {
                add_ep(right);
                erase_unassign(right);
				if (sep == len - 1 || (best_situation & (1L << sep))) break;
				sep ++;
				right ++;
			}
			for (int l = left;l <= right;l++) {
				flows[l].is_scheduled = WCRT(flows[l]);
				flows[l].new_priority = prio;
			}
			prio --;
			for (int l = left;l <= right;l++) {
                add_assign(l);
			}
		}
		i = j;
    }
	// 还原优先级
	int desc_priority = 0;
	for (int i = 2;i <= flow_num;i++) {
		if (flows[i].old_priority == flows[i - 1].old_priority && flows[i].new_priority != flows[i - 1].new_priority) {
            desc_priority ++;
		}
        flows[i].priority = flows[i].priority - desc_priority;
	}
    LPR();
}

void LPAS::LPR() {
    sort(flows + 1, flows + flow_num + 1, [&](flow& a, flow& b) {
        if (a.priority == b.priority) {
            return a.is_scheduled > b.is_scheduled;
        }
        return a.priority < b.priority;
    });
    for (int vm = 0; vm < VM_NUM; ++ vm) {
        qid2fid.clear();
        queue_load.assign(flow_num + 1, 0);
        int j = 1, last, start = 0;
        for (int i = 1;i <= flow_num; i++) {
            if (flows[i].src == vm) {
                start = last = i;
                break;
            }
        }
        for (int i = start; i <= flow_num; i++) {
            if (flows[i].src != vm) continue;
            // if (flows[i].priority != flows[i - 1].priority || flows[i].is_scheduled == false) {
            if (flows[i].priority != flows[last].priority) {
                j++;
            }
            if(qid2fid.count(j) == 0)
                qid2fid[j] = vector<int>();
            qid2fid[j].push_back(i);
            queue_load[j] += (HP[vm] / flows[i].period) * flows[i].size;
            last = i;
            // INFO("qid: %d, j: %d , queue_load: %d\n", i, j, queue_load[j]);
        }
        queue_num = j;
        average_load = accumulate(queue_load.begin(), queue_load.end(), 0) / NEW_PRIORITY_NUM;
        INFO("average_load: %llu, hp: %llu, qnum: %d\n", average_load, HP[vm], j);
        min_cost  = __LONG_LONG_MAX__;
        map_queue.resize(NEW_PRIORITY_NUM);
        for (auto& q : map_queue) {
            q.clear();
        }
        if (start)
            Reassign_priority(vm, 1, 0, 0);
    }
}

void LPAS::Reassign_priority(int vm, int qid, int cur_lpa_prio, u64 cost) {
    // Logger::log(LOG_PATH, "qid: %d, cur_lpa_prio: %d, cost: %llu, min: %llu\n", qid, cur_lpa_prio, cost,
    // min_cost);
    if (cost >= min_cost) {
        return;
    }
    if (qid > queue_num) {
        while (cur_lpa_prio < NEW_PRIORITY_NUM) {
            cost += average_load * average_load;
            cur_lpa_prio++;
        }
        if (cost < min_cost) {
            min_cost = cost;
            for (int j = 1; j <= flow_num; j++) {
                if (flows[j].src == vm) {
                    flows[j].final_priority = 0;
                }
            }
            for (int i = 0; i < NEW_PRIORITY_NUM; i++) {
                for (auto& id : map_queue[i]) {
                    flows[id].final_priority = i + 1;
                }
            }
        }
        return;
    }
    if (cur_lpa_prio == NEW_PRIORITY_NUM) {
        return;
    }
    u64 cur_load = 0;
    while (cur_load <= average_load && qid <= queue_num) {
        assert(flows[qid].src == vm);
        cur_load += queue_load[qid];
        for (const auto& fid: qid2fid[qid]) {
            map_queue[cur_lpa_prio].push_back(fid);
        }
        qid++;
        // Logger::log(LOG_PATH,
        //             "cur_load: %llu, average_load: %llu, qid: %d, cur_lpa_prio: %d, cost: %llu, min:%llu\n ",
        //             cur_load, average_load, qid, cur_lpa_prio,
        //             cost + (cur_load - average_load) * (cur_load - average_load), min_cost);
        Reassign_priority(vm, qid, cur_lpa_prio + 1,
            cost + (cur_load - average_load) * (cur_load - average_load));
    }
    map_queue[cur_lpa_prio].clear();
}


void LPAS::debug_flows() {
    for (int i = 1; i <= flow_num; i++)
        temp[i] = flows[i];
    sort(temp + 1, temp + flow_num + 1, [&](flow& a, flow& b) { return a.id < b.id; });
    printf("qid  old_priority  new_priority  is_scheduled   wcrt\n");
    for (int i = 1; i <= flow_num; i++)
        printf("%3d %12d %12d %13d %13.2f\n", i, temp[i].old_priority, temp[i].priority, temp[i].is_scheduled,
               print_wcrt(temp[i].wcrt));

    int hi_sche = 0, lo_sche = 0;
    for (int i = 1; i <= flow_num; i++) {
        if (temp[i].is_scheduled && temp[i].critical)
            hi_sche ++;
        else if (temp[i].is_scheduled && !temp[i].critical)
            lo_sche ++;
    }
    printf("hi_sche: %d, lo_sche: %d, hi rate:%lf, lo rate:%lf\n", hi_sche, lo_sche, (double)hi_sche / flow_num * 2, (double)lo_sche / flow_num * 2);
}

void LPAS::record_flows(string path) {
    std::string base = path;
    std::string ext;
    size_t dot_pos = path.rfind('.');
    if (dot_pos != std::string::npos) {
        base = path.substr(0, dot_pos);
        ext = path.substr(dot_pos); // e.g. ".csv"
    }

    int suffix = 1;
    while (fs::exists(path)) {
        path = base + "_" + std::to_string(suffix++) + ext;
    }

    std::ofstream out(path);
    if (!out.is_open()) {
        std::cerr << "Failed to open file: " << path << std::endl;
        return;
    }

    out << "id,src,dst,critical,size,period,relative_dl,priority\n";
    for (size_t i = 1; i <= flow_num; ++i) {
        out << temp[i].id << ","
            << temp[i].src << ","
            << temp[i].dst << ","
            << temp[i].critical << ","
            << temp[i].size << ","
            << temp[i].period << ","
            << temp[i].relative_dl << ","
            << temp[i].final_priority << "\n";
    }
    out.close();
}

void LPAS::continuous_record(string path) {
    std::ofstream out(path, std::ios::app);
    if (!out.is_open()) {
        std::cerr << "Failed to open file: " << path << std::endl;
        return;
    }

    out << "id,src,dst,critical,size,period,relative_dl,priority\n";
    for (size_t i = 1; i <= flow_num; ++i) {
        out << temp[i].id << ","
            << temp[i].src << ","
            << temp[i].dst << ","
            << temp[i].critical << ","
            << temp[i].size << ","
            << temp[i].period << ","
            << temp[i].relative_dl << ","
            << temp[i].final_priority << "\n";
    }
    out.close();
}